/** @file

Read macOS .disk_label (.disk_label_2x) file and convert to .ppm.
Write macOS .disk_label (.disk_label_2x) files from text string.
Reference:
- http://refit.sourceforge.net/info/vollabel.html
- https://opensource.apple.com/source/bless/bless-181.0.1/libbless/Misc/BLGenerateOFLabel.c.auto.html

Copyright (c) 2019-2020, vit9696

All rights reserved.

This program and the accompanying materials
are licensed and made available under the terms and conditions of the BSD License
which accompanies this distribution.  The full text of the license may be found at
http://opensource.org/licenses/bsd-license.php

THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/
#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>

#ifdef __APPLE__
#include <ApplicationServices/ApplicationServices.h>
#include <CoreFoundation/CoreFoundation.h>
#include <CoreGraphics/CoreGraphics.h>
#endif // __APPLE__

/* Antialiasing clut. */
static const uint8_t clut[] = {
  0x00, /* 0x00 0x00 0x00 white */
  0xF6, /* 0x11 0x11 0x11 */
  0xF7, /* 0x22 0x22 0x22 */
  0x2A, /* 0x33 = 1*6^2 + 1*6 + 1 = 43 colors */
  0xF8, /* 0x44 */
  0xF9, /* 0x55 */
  0x55, /* 0x66 = 2*(36 + 6 + 1) = 86 colors */
  0xFA, /* 0x77 */
  0xFB, /* 0x88 */
  0x80, /* 0x99 = (3*43) = 129 colors*/
  0xFC, /* 0xAA */
  0xFD, /* 0xBB */
  0xAB, /* 0xCC = 4*43 = 172 colors */
  0xFE, /* 0xDD */
  0xFF, /* 0xEE */
  0xD6, /* 0xFF = 5*43 = 215 */
};

/* Reverse clut to pixels. */
static const uint8_t palette[256] = {
  [0x00] = /* 0xFF - */ 0x00,
  [0xF6] = /* 0xFF - */ 0x11,
  [0xF7] = /* 0xFF - */ 0x22,
  [0x2A] = /* 0xFF - */ 0x33,
  [0xF8] = /* 0xFF - */ 0x44,
  [0xF9] = /* 0xFF - */ 0x55,
  [0x55] = /* 0xFF - */ 0x66,
  [0xFA] = /* 0xFF - */ 0x77,
  [0xFB] = /* 0xFF - */ 0x88,
  [0x80] = /* 0xFF - */ 0x99,
  [0xFC] = /* 0xFF - */ 0xAA,
  [0xFD] = /* 0xFF - */ 0xBB,
  [0xAB] = /* 0xFF - */ 0xCC,
  [0xFE] = /* 0xFF - */ 0xDD,
  [0xFF] = /* 0xFF - */ 0xEE,
  [0xD6] = /* 0xFF - */ 0xFF
};

/* To fit BootPicker. */
#define LABEL_MAX_WIDTH  340
#define LABEL_MAX_HEIGHT 12
#define LABEL_VERSION    1

#pragma pack(push, 1)
typedef struct DiskLabel_ {
  uint8_t  version;
  uint16_t width;
  uint16_t height;
  uint8_t  data[];
} DiskLabel;
#pragma pack(pop)

#ifdef __APPLE__
static CTLineRef create_text_line(CGColorSpaceRef color_space, const char *str, int scale) {
  /* White text on black background, originating from OF/EFI bitmap. */
  const CGFloat components[] = {
    (CGFloat)1.0, (CGFloat)1.0
  };

  CTFontRef font_ref          = CTFontCreateWithName(CFSTR("Helvetica"), 10.0 * scale, NULL);
  CGColorRef color            = CGColorCreate(color_space, components);
  CFMutableDictionaryRef dict = CFDictionaryCreateMutable(kCFAllocatorDefault,
    0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);

  CTLineRef result = NULL;

  if (font_ref != NULL && color != NULL && dict != NULL) {
    CFDictionarySetValue(dict, kCTForegroundColorAttributeName, color);
    CFDictionarySetValue(dict, kCTFontAttributeName, font_ref);

    CFStringRef tmp = CFStringCreateWithCString(kCFAllocatorDefault, str, kCFStringEncodingUTF8);
    if (tmp != NULL) {
      CFAttributedStringRef attr_tmp = CFAttributedStringCreate(kCFAllocatorDefault, tmp, dict);
      if (attr_tmp != NULL) {
        result = CTLineCreateWithAttributedString(attr_tmp);
        CFRelease(attr_tmp);
      }
      CFRelease(tmp);
    }
  }

  if (dict != NULL) {
    CFRelease(dict);
  }

  if (color != NULL) {
    CFRelease(color);
  }

  if (font_ref != NULL) {
    CFRelease(font_ref);
  }

  return result;
}

static int render_label(const char *label, uint8_t *bitmap_data,
  uint16_t width, uint16_t height, int scale, uint16_t *newwidth) {

  CGColorSpaceRef color_space = CGColorSpaceCreateDeviceGray();
  if (color_space == NULL) {
    fprintf(stderr, "Could not obtain color space\n");
    return 1;
  }

  CTLineRef line = create_text_line(color_space, label, scale);
  if (line == NULL) {
    fprintf(stderr, "Could not render text line\n");
    CGColorSpaceRelease(color_space);
    return 3;
  }

  CGContextRef context = CGBitmapContextCreate(bitmap_data, width, height,
    8, width, color_space, kCGImageAlphaNone);

  if(context == NULL) {
    fprintf(stderr, "Could not init CoreGraphics context\n");
    CFRelease(line);
    CGColorSpaceRelease(color_space);
    return 2;
  }

  CGRect rect = CTLineGetImageBounds(line, context);
  CGContextSetTextPosition(context, 2.0 * scale, 2.0 * scale);
  CTLineDraw(line, context);
  CGContextFlush(context);
  CFRelease(line);

  *newwidth = (uint16_t)(rect.size.width + 4 * scale);

  CGContextRelease(context);
  CGColorSpaceRelease(color_space);
  return 0;
}

static void *make_label(const char *label, int scale, uint16_t *label_size) {
  uint16_t width  = LABEL_MAX_WIDTH * scale;
  uint16_t height = LABEL_MAX_HEIGHT * scale;
  DiskLabel *label_data = calloc(1, sizeof(DiskLabel) + width*height);
  if (label_data == NULL) {
    fprintf(stderr, "Label allocation failure\n");
    return NULL;
  }

  uint16_t newwidth;
  int err = render_label(label, label_data->data, width, height, scale, &newwidth);
  if (err != 0) {
    free(label_data);
    return NULL;
  }

  /* Cap at 340*scale pixels wide. */
  if (newwidth > width) {
    fprintf(stderr, "Label %s with %d pixels does not fit in %d pixels\n", label, newwidth, width);
    newwidth = width;
  }

  /* Refit to new width (111111000111111000111111000 -> 111111111111111111). */
  for (uint16_t row = 1; row < height; row++) {
    memmove(&label_data->data[row * newwidth], &label_data->data[row * width], newwidth);
  }

  label_data->version = LABEL_VERSION;
  label_data->width   = htons(newwidth);
  label_data->height  = htons(height);

  /* Perform antialiasing. */
  for (uint16_t i = 0; i < newwidth * height; i++) {
    label_data->data[i] = clut[label_data->data[i] >> 4U];
  }

  *label_size = sizeof(DiskLabel) + newwidth * height;

  return label_data;
}
#endif // __APPLE__

static int read_file(const char *filename, uint8_t **buffer, size_t *size) {
  FILE *fh = fopen(filename, "rb");
  if (!fh) {
    fprintf(stderr, "Missing file %s!\n", filename);
    return -1;
  }

  if (fseek(fh, 0, SEEK_END)) {
    fprintf(stderr, "Failed to find end of %s!\n", filename);
    fclose(fh);
    return -1;
  }

  long pos = ftell(fh);

  if (pos <= 0) {
    fprintf(stderr, "Invalid file size (%ld) of %s!\n", pos, filename);
    fclose(fh);
    return -1;
  }

  if (fseek(fh, 0, SEEK_SET)) {
    fprintf(stderr, "Failed to rewind %s!\n", filename);
    fclose(fh);
    return -1;
  }

  *size = (size_t)pos;
  *buffer = (uint8_t *)malloc(*size);

  if (!*buffer) {
    fprintf(stderr, "Failed to allocate %zu bytes for %s!\n", *size, filename);
    fclose(fh);
    return -1;
  }

  if (fread(*buffer, *size, 1, fh) != 1) {
    fprintf(stderr, "Failed to read %zu bytes from %s!\n", *size, filename);
    fclose(fh);
    free(*buffer);
    return -1;
  }

  fclose(fh);
  return 0;
}

static int write_file(const char *filename, uint8_t *buffer, size_t size) {
  FILE *fh = fopen(filename, "wb+");
  if (!fh) {
    fprintf(stderr, "Cannot open file %s for writing!\n", filename);
    return -1;
  }

  if (fwrite(buffer, size, 1, fh) != 1) {
    fprintf(stderr, "Cannot write %zu bytes in %s!\n", size, filename);
    fclose(fh);
    return -1;
  }

  fclose(fh);
  return 0;
}

static int write_ppm(const char *filename, size_t width, size_t height, uint8_t *pixel) {
  FILE *fh = fopen(filename, "wb");
  if (!fh) {
    fprintf(stderr, "Failed to open out file %s!\n", filename);
    return -1;
  }

  if (fprintf(fh, "P6\n%zu %zu\n255\n", width, height) < 0) {
    fprintf(stderr, "Failed to write ppm header to %s!\n", filename);
    fclose(fh);
    return -1;
  }

  for (size_t y = 0; y < height; ++y) {
    for (size_t x = 0; x < width; ++x) {
      uint8_t col[3];
      col[0] = col[1] = col[2] = palette[*pixel];
      if (fwrite(col, sizeof(col), 1, fh) != 1) {
        fprintf(stderr, "Failed to write ppm pixel %zux%zu to %s!\n", x, y, filename);
        fclose(fh);
        return -1;
      }
      ++pixel;
    }
  }

  fclose(fh);

  return 0;
}

static int decode_label(const char *infile, const char *outfile) {
  DiskLabel *label;
  size_t size;

  if (read_file(infile, (uint8_t **) &label, &size)) {
    return -1;
  }

  if (size < sizeof(DiskLabel)) {
    fprintf(stderr, "Too low size %zu!\n", size);
    free(label);
    return -1;
  }

  if (label->version != LABEL_VERSION) {
    fprintf(stderr, "Invalid version %02X!\n", label->version);
    free(label);
    return -1;
  }

  size_t width  = ntohs(label->width);
  size_t height = ntohs(label->height);

  if (width * height == 0 || size - sizeof(DiskLabel) != width*height) {
    fprintf(stderr, "Image mismatch %zux%zu with size %zu!\n", width, height, size);
    free(label);
    return -1;
  }

  (void)remove(outfile);

  int ret = write_ppm(outfile, width, height, label->data);
  free(label);
  return ret;
}

static int encode_label(const char *label, const char *outfile, const char *outfile2x) {
#ifdef __APPLE__
  for (int scale = 1; scale <= 2; scale++) {
    const char *filename = scale == 1 ? outfile : outfile2x;
    uint16_t label_size;
    void *label_data = make_label(label, scale, &label_size);
    bool ok = label_data != NULL && write_file(filename, label_data, label_size) == 0;
    free(label_data);
    if (!ok) {
      return -2;
    }
  }
  return 0;
#else
  (void) label;
  (void) outfile;
  (void) outfile2x;
  fprintf(stderr, "Encoding labels is unsupported!\n");
  return -1;
#endif
}

int main(int argc, char *argv[]) {
  if (argc == 4 && strcmp(argv[1], "-d") == 0) {
    return decode_label(argv[2], argv[3]);
  }

  if (argc == 5 && strcmp(argv[1], "-e") == 0) {
    return encode_label(argv[2], argv[3], argv[4]);
  }

  fprintf(stderr,
    "Usage:\n"
    " disklabel -d .disk_label image.ppm!\n"
    " disklabel -e \"Label\" .disk_label .disk_label_2x\n");
  return -1;
}
